#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
修复三班制排班接口的日期处理问题
"""

import logging
from datetime import datetime, date
from app.models import ScheduleResult, TimeSlot, ShiftAssignment
from app import db

# 设置日志
logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.INFO)

def fix_date_issues():
    """修复三班制排班接口的日期处理问题"""
    logger.info("开始修复三班制排班接口的日期处理问题")
    
    try:
        # 查询所有plan_date为2025-09-09, 2025-09-10, 2025-09-11的排班结果
        problematic_dates = ['2025-09-09', '2025-09-10', '2025-09-11']
        results_to_fix = []
        
        for date_str in problematic_dates:
            date_obj = datetime.strptime(date_str, '%Y-%m-%d').date()
            results = ScheduleResult.query.filter(ScheduleResult.plan_date == date_obj).all()
            results_to_fix.extend(results)
        
        logger.info(f"找到 {len(results_to_fix)} 条需要修复的排班结果记录")
        
        # 按批次号分组
        batch_groups = {}
        for result in results_to_fix:
            if result.batch_number not in batch_groups:
                batch_groups[result.batch_number] = []
            batch_groups[result.batch_number].append(result)
        
        logger.info(f"共有 {len(batch_groups)} 个不同的批次")
        
        # 为每个批次分配正确的日期
        for batch_number, results in batch_groups.items():
            # 查找该批次的班次分配记录，获取正确的日期
            shift_assignment = ShiftAssignment.query.filter(
                ShiftAssignment.created_at >= results[0].schedule_date - datetime.timedelta(minutes=5),
                ShiftAssignment.created_at <= results[0].schedule_date + datetime.timedelta(minutes=5)
            ).first()
            
            if shift_assignment:
                correct_date = shift_assignment.shift_date
                logger.info(f"批次 {batch_number} 的正确日期应为: {correct_date}")
                
                # 更新排班结果的plan_date
                for result in results:
                    result.plan_date = correct_date
                    logger.info(f"更新排班结果 {result.id} 的plan_date为 {correct_date}")
                
                # 更新相关的时间段记录
                for result in results:
                    time_slots = TimeSlot.query.filter_by(schedule_result_id=result.id).all()
                    for slot in time_slots:
                        slot.plan_date = correct_date
                    
                    logger.info(f"更新了 {len(time_slots)} 条时间段记录的plan_date")
            else:
                # 如果找不到对应的班次分配记录，使用schedule_date的日期部分
                for result in results:
                    correct_date = result.schedule_date.date()
                    result.plan_date = correct_date
                    logger.info(f"无法找到班次分配记录，使用schedule_date的日期部分: {correct_date}")
                    
                    # 更新相关的时间段记录
                    time_slots = TimeSlot.query.filter_by(schedule_result_id=result.id).all()
                    for slot in time_slots:
                        slot.plan_date = correct_date
                    
                    logger.info(f"更新了 {len(time_slots)} 条时间段记录的plan_date")
        
        # 提交更改
        db.session.commit()
        logger.info("成功提交所有更改")
        return True
    except Exception as e:
        db.session.rollback()
        logger.error(f"修复日期问题失败: {str(e)}")
        return False

if __name__ == "__main__":
    # 直接运行此脚本时，执行修复操作
    from app import create_app
    with create_app().app_context():
        fix_date_issues()
